import os
import numpy as np
from PIL import Image
from scipy.ndimage import gaussian_filter


class TransparentBgProcessor:
    """
    图片透明背景处理工具包
    可用于去除指定颜色（如白色背景），添加羽化效果，并裁剪多余的透明区域
    """

    @staticmethod
    def remove_color(
            image_path=None,
            image=None,
            target_color=(255, 255, 255),
            tolerance=0,
            blur_radius=2,
    ):
        """
        读取图片并去除指定颜色的像素，将目标颜色像素变成透明，并添加羽化效果

        参数:
            image_path: 图片文件路径，与image参数二选一
            image: PIL图片对象，与image_path参数二选一
            target_color: 要去除的目标RGB颜色，默认为白色 (255, 255, 255)
            tolerance: 颜色匹配的容差值，0表示精确匹配，值越大匹配范围越广
            blur_radius: 羽化半径，值越大羽化效果越明显

        返回:
            处理后的PIL图片对象
        """
        try:
            # 检查输入参数
            if image_path is None and image is None:
                raise ValueError("必须提供图片路径或图片对象")

            # 打开图片
            if image is None:
                img = Image.open(image_path)
            else:
                img = image

            # 确保图片是RGBA模式
            if img.mode != "RGBA":
                img = img.convert("RGBA")

            # 将图片转换为numpy数组以便处理
            img_array = np.array(img)

            # 解包目标颜色
            r_target, g_target, b_target = target_color

            # 创建掩码，找出符合目标颜色的像素（考虑容差值）
            if tolerance == 0:
                # 精确匹配
                target_pixels = (
                        (img_array[:, :, 0] == r_target)
                        & (img_array[:, :, 1] == g_target)
                        & (img_array[:, :, 2] == b_target)
                )
            else:
                # 带容差的匹配
                target_pixels = (
                        (
                                np.abs(img_array[:, :, 0].astype(np.int16) - r_target)
                                <= tolerance
                        )
                        & (
                                np.abs(img_array[:, :, 1].astype(np.int16) - g_target)
                                <= tolerance
                        )
                        & (
                                np.abs(img_array[:, :, 2].astype(np.int16) - b_target)
                                <= tolerance
                        )
                )

            # 创建新的数组
            new_array = img_array.copy()

            # 将目标颜色像素的alpha通道设置为0（完全透明）
            new_array[target_pixels, 3] = 0

            # 创建alpha通道的掩码
            alpha_mask = new_array[:, :, 3].copy()

            # 对alpha通道应用高斯模糊，创建羽化效果
            blurred_alpha = gaussian_filter(alpha_mask, sigma=blur_radius)

            # 将模糊后的alpha通道应用到图片
            new_array[:, :, 3] = blurred_alpha

            # 创建新的RGBA图片
            result = Image.fromarray(new_array.astype(np.uint8))

            return result

        except Exception as e:
            print(f"处理图片时发生错误: {str(e)}")
            return None

    @staticmethod
    def crop_transparent_margins(image, margin=0):
        """
        裁切图片的透明边缘，只保留非透明区域加上指定的边距

        参数:
            image: PIL图片对象（RGBA模式）
            margin: 保留的边距像素数

        返回:
            裁切后的PIL图片对象
        """
        try:
            # 确保图片是RGBA模式
            if image.mode != "RGBA":
                image = image.convert("RGBA")

            # 获取图片数组
            img_array = np.array(image)

            # 获取非透明像素的位置（alpha通道值大于0的像素）
            non_transparent = img_array[:, :, 3] > 0

            # 如果图片完全透明，则返回原图
            if not np.any(non_transparent):
                return image

            # 获取非透明区域的边界
            rows = np.any(non_transparent, axis=1)
            cols = np.any(non_transparent, axis=0)

            # 找到非透明区域的最小和最大行列索引
            min_row, max_row = np.where(rows)[0][[0, -1]]
            min_col, max_col = np.where(cols)[0][[0, -1]]

            # 应用边距，并确保不超出图片边界
            min_row = max(0, min_row - margin)
            min_col = max(0, min_col - margin)
            max_row = min(img_array.shape[0] - 1, max_row + margin)
            max_col = min(img_array.shape[1] - 1, max_col + margin)

            # 裁切图片
            return image.crop((min_col, min_row, max_col + 1, max_row + 1))

        except Exception as e:
            print(f"裁切图片时发生错误: {str(e)}")
            return image

    @staticmethod
    def save_image(image, output_path):
        """
        保存处理后的图片

        参数:
            image: 处理后的图片对象
            output_path: 输出文件路径
        """
        try:
            # 确保输出目录存在
            output_dir = os.path.dirname(output_path)
            if output_dir and not os.path.exists(output_dir):
                os.makedirs(output_dir)

            # 保存图片，确保使用PNG格式以支持透明度
            image.save(output_path, "PNG")
            print(f"图片已成功保存到: {output_path}")
            return True
        except Exception as e:
            print(f"保存图片时发生错误: {str(e)}")
            return False

    @classmethod
    def process_image(
            cls,
            input_path=None,
            input_image=None,
            target_color=(255, 255, 255),
            tolerance=0,
            blur_radius=2,
            margin=0,
    ):
        """
        处理图片的主函数

        参数:
            input_path: 输入图片路径，与input_image二选一
            input_image: 输入图片对象，与input_path二选一
            target_color: 要去除的目标RGB颜色，默认为白色 (255, 255, 255)
            tolerance: 颜色匹配的容差值，0表示精确匹配，值越大匹配范围越广
            blur_radius: 羽化半径，值越大羽化效果越明显
            margin: 裁切时保留的透明边距像素数

        返回:
            处理后的PIL图片对象，处理失败则返回None
        """
        try:
            # 检查输入参数
            if input_path is None and input_image is None:
                raise ValueError("必须提供输入图片路径或图片对象")

            # 检查输入文件是否存在
            if input_path is not None and not os.path.exists(input_path):
                print(f"输入文件不存在: {input_path}")
                return None

            # 处理图片
            processed_image = cls.remove_color(
                image_path=input_path,
                image=input_image,
                target_color=target_color,
                tolerance=tolerance,
                blur_radius=blur_radius,
            )

            if processed_image:
                # 裁切透明边缘
                if margin >= 0:  # 仅在margin大于等于0时进行裁切
                    processed_image = cls.crop_transparent_margins(
                        processed_image, margin
                    )

                return processed_image

            return None

        except Exception as e:
            print(f"处理图片时发生错误: {str(e)}")
            return None


# 使用示例
if __name__ == "__main__":
    # 示例用法
    input_image_path = (
        r"N:\trainData\衣服\ComfyUI_temp_eruei_00003_.png"  # 替换为你的输入图片路径
    )
    output_image_path = r"N:\trainData\衣服\output.png"  # 替换为你想保存的输出图片路径

    img = Image.open(input_image_path)

    # 示例1：移除白色像素 (255, 255, 255)
    processed_image = TransparentBgProcessor.process_image(
        input_image=img,
        target_color=(255, 255, 255),  # 白色
        tolerance=5,  # 允许5个单位的颜色误差
        blur_radius=5,
        margin=10,
    )

    # 保存处理后的图片
    if processed_image:
        TransparentBgProcessor.save_image(processed_image, output_image_path)

    # # 示例2：从PIL图像对象处理和保存
    # from PIL import Image
    # original_image = Image.open(input_image_path)
    # processed_image = TransparentBgProcessor.process_image(
    #     input_image=original_image,
    #     target_color=(0, 0, 0),  # 黑色
    #     tolerance=10,  # 允许10个单位的颜色误差
    #     blur_radius=3,
    #     margin=5,
    # )
    # if processed_image:
    #     TransparentBgProcessor.save_image(processed_image, "output_black_removed.png")
